'''
优惠政策，售后政策
'''
import re
import openai
from langchain_core.tools import tool
import os
from dotenv import load_dotenv
from langchain_chroma.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from chromadb.config import Settings
from tools.utility import timer_decorator
load_dotenv()


def load_policy():
# 读取policy文件夹中的所有文件
    policy_dir = os.path.join(os.path.dirname(__file__), "policy")
    faq_text = ""

    # 遍历policy目录下的所有文件
    docs = [TextLoader(os.path.join(policy_dir, filename)).load() for filename in os.listdir(policy_dir)]
    docs_list = [item for sublist in docs for item in sublist]

    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        chunk_size=100, chunk_overlap=50
    )
    doc_splits = text_splitter.split_documents(docs_list)


    vector_store = Chroma(
        collection_name="policies",
        embedding_function=OpenAIEmbeddings(),
        client_settings=Settings(allow_reset=True)
    )
    vector_store.reset_collection()
    vector_store.add_documents(doc_splits)

    retriever = vector_store.as_retriever()
@timer_decorator
def get_retriever():
    vector_store = Chroma(
        collection_name="policies",
        embedding_function=OpenAIEmbeddings(),
        client_settings=Settings(allow_reset=True)
    )
    retriever = vector_store.as_retriever()
    return retriever


from langchain.tools.retriever import create_retriever_tool






@tool
@timer_decorator
def lookup_policy(query: str) -> str:
    """
    查询平台优惠政策和退换货政策

    Args:
        query (str): 查询内容

    Returns:
        str: 查询结果
    """
    retriever = get_retriever()

    return retriever.invoke(query)





if __name__ == "__main__":

    print(lookup_policy.invoke("优惠券"))
